#include "stdafx.h"

#include "DetouredFunctions.hpp"
#include "OutputConfig.h"
#include "DataCollectionTypes.h"
#include "PostOffice.h"
#include "WinsockHooks.h"

#include <detours.h>
#include <process.h>
#include <fstream>
#include <tlhelp32.h>


// Target pointer for the uninstrumented API functions.
WINAPI_CreateFileW OriginalAPI::CreateFileW = ::CreateFileW;
WINAPI_CreateFileA OriginalAPI::CreateFileA = ::CreateFileA;
WINAPI_WriteFile OriginalAPI::WriteFile = ::WriteFile;
WINAPI_ReadFile OriginalAPI::ReadFile = ::ReadFile;

WINAPI_VirtualAlloc OriginalAPI::VirtualAlloc = ::VirtualAlloc;
WINAPI_VirtualAllocEx OriginalAPI::VirtualAllocEx = ::VirtualAllocEx;

WINAPI_ExitProcess OriginalAPI::ExitProcess = ::ExitProcess;
WINAPI_DuplicateHandle OriginalAPI::DuplicateHandle = ::DuplicateHandle;

WINAPI_CreateThread OriginalAPI::CreateThread = ::CreateThread;
WINAPI_ExitThread OriginalAPI::ExitThread = ::ExitThread;
WINAPI_TerminateThread OriginalAPI::TerminateThread = ::TerminateThread;

WINAPI_CreateWindowExA OriginalAPI::CreateWindowExA = ::CreateWindowExA;
WINAPI_CreateWindowExW OriginalAPI::CreateWindowExW = ::CreateWindowExW;

WINAPI_LoadLibraryA OriginalAPI::LoadLibraryA = ::LoadLibraryA;
WINAPI_LoadLibraryW OriginalAPI::LoadLibraryW = ::LoadLibraryW;
WINAPI_LoadLibraryExA OriginalAPI::LoadLibraryExA = ::LoadLibraryExA;
WINAPI_LoadLibraryExW OriginalAPI::LoadLibraryExW = ::LoadLibraryExW;

WINAPI_WaitForSingleObject OriginalAPI::WaitForSingleObject = ::WaitForSingleObject;
WINAPI_WaitForSingleObjectEx OriginalAPI::WaitForSingleObjectEx = ::WaitForSingleObjectEx;
WINAPI_WaitForMultipleObjects OriginalAPI::WaitForMultipleObjects = ::WaitForMultipleObjects;
WINAPI_WaitForMultipleObjectsEx OriginalAPI::WaitForMultipleObjectsEx = ::WaitForMultipleObjectsEx;
WINAPI_MsgWaitForMultipleObjects OriginalAPI::MsgWaitForMultipleObjects = ::MsgWaitForMultipleObjects;
WINAPI_MsgWaitForMultipleObjectsEx OriginalAPI::MsgWaitForMultipleObjectsEx = ::MsgWaitForMultipleObjectsEx;



// flags to indicate what type of hooks were installed
static bool hasInstalledFileEvents = false;
static bool hasInstalledMemoryEvents = false;
static bool hasInstalledThreadEvents = false;
static bool hasInstalledWindowEvents = false;
static bool hasInstalledLoadLibraryHooks = false;
static bool hasInstalledSynchronizationHooks = false;

static uintptr_t po_thread = 0;

static DWORD proc_start_time = 0;


/**
 * This function is called once a new thread was created for
 * PostOffice, and it will be used to write all events to
 * the output file
 */
unsigned int WINAPI InitializePostOfficeInThread(void*) {
	StateIndicator::Setter set_internal;

	config::OutputConfig conf;
	platform::XString fname = conf.getProcessOutputFile(_T("abpa"));

	std::ofstream out(fname.c_str(), std::ios_base::binary);
	if (!out.is_open())
		return 1;

	events::PostOffice::GetInstance().startLoggerInThisThread(out, proc_start_time);

	return 0;
}


#define MAX_NUMBER_OF_THREADS 10
unsigned int getListOfThreads(HANDLE *pOut) {
	// get snapshot of all threads in the system
	// NOTE: the 2nd parameter (process id) is ignored, so
	// we have to check the pid for each process
	HANDLE h = ::CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, 0);
	if (h == INVALID_HANDLE_VALUE)
		return 0;

	// cache ID of current process
	DWORD myProcId = ::GetCurrentProcessId();
	DWORD myThreadId = ::GetCurrentThreadId();

	// prepare the struct to read threads' data
	THREADENTRY32 te;
	te.dwSize = sizeof(te);

	// iterate over the list
	unsigned int idx = 0;
	if (::Thread32First(h, &te)) {
		do {
			if (te.th32OwnerProcessID == myProcId && te.th32ThreadID != myThreadId) {
				pOut[idx] = ::OpenThread(THREAD_ALL_ACCESS, FALSE, te.th32ThreadID);
				if (pOut[idx] != NULL)
					idx++;
			}
		} while (idx < MAX_NUMBER_OF_THREADS && ::Thread32Next(h, &te));
	}

	return idx;
}

/**
 * Installs all Windows API hooks using "Detours" library
 */
void InstallHooks(void) {
	StateIndicator::Setter set_internal;
	proc_start_time = ::GetTickCount();

	// get list og all threads (excluding this one)
	HANDLE hThreads[MAX_NUMBER_OF_THREADS];
	unsigned int iNumThreads = getListOfThreads(hThreads);

	// initialize PostOffice
	unsigned int thId;
	po_thread = _beginthreadex(NULL, 0, InitializePostOfficeInThread, NULL, 0, &thId);

	// install hooks
	DetourTransactionBegin();
	DetourUpdateThread(::GetCurrentThread());
	for (unsigned int i=0; i < iNumThreads; i++) DetourUpdateThread(hThreads[i]);

	config::DataCollectionTypes dtc;

	// file-related functions
	if (dtc.shouldCollectFileEvents()) {
		hasInstalledFileEvents = true;
		DetourAttach(&(PVOID&)OriginalAPI::CreateFileW, HookCreateFileW);
		DetourAttach(&(PVOID&)OriginalAPI::CreateFileA, HookCreateFileA);
		DetourAttach(&(PVOID&)OriginalAPI::WriteFile, HookWriteFile);
		DetourAttach(&(PVOID&)OriginalAPI::ReadFile, HookReadFile);
	}

	if (dtc.shouldCollectMemoryEvents()) {
		hasInstalledMemoryEvents = true;
		DetourAttach(&(PVOID&)OriginalAPI::VirtualAlloc, HookVirtualAlloc);
		DetourAttach(&(PVOID&)OriginalAPI::VirtualAllocEx, HookVirtualAllocEx);
	}

	if (dtc.shouldCollectThreadEvents()) {
		hasInstalledThreadEvents = true;
		DetourAttach(&(PVOID&)OriginalAPI::CreateThread, HookCreateThread);
		DetourAttach(&(PVOID&)OriginalAPI::ExitThread, HookExitThread);
		DetourAttach(&(PVOID&)OriginalAPI::TerminateThread, HookTerminateThread);
	}

	if (dtc.shouldCollectWindowEvents()) {
		hasInstalledWindowEvents = true;
		DetourAttach(&(PVOID&)OriginalAPI::CreateWindowExA, HookCreateWindowExA);
		DetourAttach(&(PVOID&)OriginalAPI::CreateWindowExW, HookCreateWindowExW);
	}

	if (dtc.shouldCollectNetworkEvents()) {
		hasInstalledLoadLibraryHooks = true;
		WinsockHooks::GetInstance().activate();
		DetourAttach(&(PVOID&)OriginalAPI::LoadLibraryA, HookLoadLibraryA);
		DetourAttach(&(PVOID&)OriginalAPI::LoadLibraryW, HookLoadLibraryW);
		DetourAttach(&(PVOID&)OriginalAPI::LoadLibraryExA, HookLoadLibraryExA);
		DetourAttach(&(PVOID&)OriginalAPI::LoadLibraryExW, HookLoadLibraryExW);
	}

	if (dtc.shouldCollectSyncronizationEvents()) {
		hasInstalledSynchronizationHooks = true;
		DetourAttach(&(PVOID&)OriginalAPI::WaitForSingleObject, HookWaitForSingleObject);
		DetourAttach(&(PVOID&)OriginalAPI::WaitForSingleObjectEx, HookWaitForSingleObjectEx);
		DetourAttach(&(PVOID&)OriginalAPI::WaitForMultipleObjects, HookWaitForMultipleObjects);
		DetourAttach(&(PVOID&)OriginalAPI::WaitForMultipleObjectsEx, HookWaitForMultipleObjectsEx);
		DetourAttach(&(PVOID&)OriginalAPI::MsgWaitForMultipleObjects, HookMsgWaitForMultipleObjects);
		DetourAttach(&(PVOID&)OriginalAPI::MsgWaitForMultipleObjectsEx, HookMsgWaitForMultipleObjectsEx);
	}

	// general functions
	DetourAttach(&(PVOID&)OriginalAPI::DuplicateHandle, HookDuplicateHandle);
	DetourAttach(&(PVOID&)OriginalAPI::ExitProcess, HookExitProcess);

	DetourTransactionCommit();

	for (unsigned int i=0; i < iNumThreads; i++) ::CloseHandle(hThreads[i]);
}


/**
 * Restores original functions' pointers
 */
void RestoreHooks(void) {
	DetourTransactionBegin();
	DetourUpdateThread(::GetCurrentThread());

	// file-related functions
	if (hasInstalledFileEvents) {
		DetourDetach(&(PVOID&)OriginalAPI::CreateFileW, HookCreateFileW);
		DetourDetach(&(PVOID&)OriginalAPI::CreateFileA, HookCreateFileA);
		DetourDetach(&(PVOID&)OriginalAPI::WriteFile, HookWriteFile);
		DetourDetach(&(PVOID&)OriginalAPI::ReadFile, HookReadFile);
	}

	if (hasInstalledMemoryEvents) {
		DetourDetach(&(PVOID&)OriginalAPI::VirtualAlloc, HookVirtualAlloc);
		DetourDetach(&(PVOID&)OriginalAPI::VirtualAllocEx, HookVirtualAllocEx);
	}

	if (hasInstalledThreadEvents) {
		DetourDetach(&(PVOID&)OriginalAPI::CreateThread, HookCreateThread);
		DetourDetach(&(PVOID&)OriginalAPI::ExitThread, HookExitThread);
		DetourDetach(&(PVOID&)OriginalAPI::TerminateThread, HookTerminateThread);
	}

	if (hasInstalledWindowEvents ) {
		DetourDetach(&(PVOID&)OriginalAPI::CreateWindowExA, HookCreateWindowExA);
		DetourDetach(&(PVOID&)OriginalAPI::CreateWindowExW, HookCreateWindowExW);
	}

	if (hasInstalledLoadLibraryHooks) {
		DetourDetach(&(PVOID&)OriginalAPI::LoadLibraryA, HookLoadLibraryA);
		DetourDetach(&(PVOID&)OriginalAPI::LoadLibraryW, HookLoadLibraryW);
		DetourDetach(&(PVOID&)OriginalAPI::LoadLibraryExA, HookLoadLibraryExA);
		DetourDetach(&(PVOID&)OriginalAPI::LoadLibraryExW, HookLoadLibraryExW);
	}

	if (hasInstalledSynchronizationHooks) {
		DetourDetach(&(PVOID&)OriginalAPI::WaitForSingleObject, HookWaitForSingleObject);
		DetourDetach(&(PVOID&)OriginalAPI::WaitForSingleObjectEx, HookWaitForSingleObjectEx);
		DetourDetach(&(PVOID&)OriginalAPI::WaitForMultipleObjects, HookWaitForMultipleObjects);
		DetourDetach(&(PVOID&)OriginalAPI::WaitForMultipleObjectsEx, HookWaitForMultipleObjectsEx);
		DetourDetach(&(PVOID&)OriginalAPI::MsgWaitForMultipleObjects, HookMsgWaitForMultipleObjects);
		DetourDetach(&(PVOID&)OriginalAPI::MsgWaitForMultipleObjectsEx, HookMsgWaitForMultipleObjectsEx);
	}

	// general functions
	DetourDetach(&(PVOID&)OriginalAPI::DuplicateHandle, HookDuplicateHandle);
	DetourDetach(&(PVOID&)OriginalAPI::ExitProcess, HookExitProcess);

	DetourTransactionCommit();
}

